#include "Animation.h"
#include "Scene.h"

using namespace OBR;

Animation::Animation()
{
}

Animation::~Animation()
{
}

void Animation::updateLocalToWorld(Matrix4x4f l2w)
{
    // NOTE: do nothing here because moving surface lights is not supported
}

void Animation::linkChild(VirtualProgramManager materials, VirtualProgramManager mediums,
                          VirtualProgramManager meshes, VirtualProgramManager lights,
                          std::vector<std::shared_ptr<Instance>> instances)
{
    this->child = instances[this->child_id];
}

MatrixAnimation::MatrixAnimation()
{
}

MatrixAnimation::~MatrixAnimation()
{
}

void MatrixAnimation::buildAccelStructure(OptixDeviceContext &ctx, CUstream &stream)
{
    /* ------------------------ check current status ------------------------ */
    if (this->built)
    {
        return;
    }

    /* ----------------------------- build child ---------------------------- */
    this->child->buildAccelStructure(ctx, stream);
    if (!this->child->built)
    {
        std::cerr << "children cannot be built" << std::endl;
        throw std::runtime_error("children cannot be built");
    }

    /* ----------------------------- build self ----------------------------- */
    uint num_keys = this->transforms.size();
    if (num_keys == 0)
    {
        this->handle = this->child->handle;
    }
    else
    {
        OptixMatrixMotionTransform *motion_transform;
        uint size;
        // get block
        this->setupMotionTransformBlock((void **)&motion_transform, num_keys, size);
        // to device
        this->d_transform_buffer.resize(size);
        this->d_transform_buffer.upload(motion_transform, size);
        // convert to handle
        OPTIX_CHECK(optixConvertPointerToTraversableHandle(
            ctx,                                            // device
            this->d_transform_buffer.d_pointer(),           // ptr
            OPTIX_TRAVERSABLE_TYPE_MATRIX_MOTION_TRANSFORM, // ptr type
            &(this->handle)                                 // dest)
            ));
        free(motion_transform);
    }

    this->built = true;
}

void MatrixAnimation::updateAccelStructure()
{
    uint num_keys = this->transforms.size();
    OptixMatrixMotionTransform *motion_transform;
    uint size;
    // get block
    this->setupMotionTransformBlock((void **)&motion_transform, num_keys, size);
    // upload
    this->d_transform_buffer.upload(motion_transform, size);
}

void MatrixAnimation::_buildFromConfig(const Json::Value &config)
{
    this->tmax = config["tmax"].asFloat();
    this->tmin = config["tmin"].asFloat();
    // TODO: support motion flag in config
    this->child_id = config["child id"].asInt(); // child will be bind later

    for (Json::ArrayIndex i = 0; i < config["transforms"].size(); i++)
    {
        if (config["transforms"][i]["type"].asString().compare("matrix") == 0)
        {
            this->transforms.push_back(ConfigHelpers::loadMatrixTransform(config["transforms"][i]));
        }
        else if (config["transforms"][i]["type"].asString().compare("srt") == 0)
        {
            this->transforms.push_back(ConfigHelpers::loadSRTTransform(config["transforms"][i]));
        }
        else
        {
            throw std::runtime_error("invalid transform type: " +
                                     config["transforms"][i]["type"].asString());
        }
    }
}

void MatrixAnimation::setupMotionTransformBlock(void **dst, uint &num_keys, uint &size)
{
    const uint matsize = sizeof(float) * 12;
    size = sizeof(OptixMatrixMotionTransform) + max(0, num_keys - 2) * matsize;
    // init transform block
    *dst = malloc(size);
    memset(*dst, 0, size);
    // set values
    OptixMatrixMotionTransform *motion_transform =
        reinterpret_cast<OptixMatrixMotionTransform *>(*dst);
    motion_transform->child = this->child->handle;
    motion_transform->motionOptions.numKeys = max(2, num_keys);
    motion_transform->motionOptions.flags = this->flags;
    motion_transform->motionOptions.timeBegin = this->tmin;
    motion_transform->motionOptions.timeEnd = this->tmax;
    // copy matrix
    float *p_transform = (float *)(motion_transform->transform);
    uint offset = 0;
    for (auto &t : this->transforms)
    {
        Matrix4x4f mat = t->getTransformMatrix();
        memcpy(p_transform + offset, mat.data, matsize);
        offset += 12;
    }
    if (num_keys == 1)
    {
        Matrix4x4f mat = this->transforms[0]->getTransformMatrix();
        memcpy(p_transform + 12, mat.data, matsize);
    }
}

SRTAnimation::SRTAnimation()
{
}

SRTAnimation::~SRTAnimation()
{
}

void SRTAnimation::buildAccelStructure(OptixDeviceContext &ctx, CUstream &stream)
{
    /* ------------------------ check current status ------------------------ */
    if (this->built)
    {
        return;
    }

    /* ----------------------------- build child ---------------------------- */
    child->buildAccelStructure(ctx, stream);
    if (!this->child->built)
    {
        std::cerr << "children cannot be built" << std::endl;
        throw std::runtime_error("children cannot be built");
    }

    /* ----------------------------- build self ----------------------------- */
    uint num_keys = this->transforms.size();
    if (num_keys == 0)
    {
        this->handle = this->child->handle;
    }
    else
    {
        OptixSRTMotionTransform *motion_transform;
        uint size;
        // get block
        this->setupMotionTransformBlock((void **)&motion_transform, num_keys, size);
        // to device
        this->d_transform_buffer.resize(size);
        this->d_transform_buffer.upload(motion_transform, size);
        // convert to handle
        OPTIX_CHECK(optixConvertPointerToTraversableHandle(
            ctx,                                         // device
            this->d_transform_buffer.d_pointer(),        // ptr
            OPTIX_TRAVERSABLE_TYPE_SRT_MOTION_TRANSFORM, // ptr type
            &(this->handle)                              // dest
            ));
        free(motion_transform);
    }

    this->built = true;
}

void SRTAnimation::updateAccelStructure()
{
    uint num_keys = this->transforms.size();
    OptixSRTMotionTransform *motion_transform;
    uint size;
    // get block
    this->setupMotionTransformBlock((void **)&motion_transform, num_keys, size);
    // upload
    this->d_transform_buffer.upload(motion_transform, size);
}

void SRTAnimation::_buildFromConfig(const Json::Value &config)
{
    this->tmax = config["tmax"].asFloat();
    this->tmin = config["tmin"].asFloat();
    // TODO: support motion flag in config
    this->child_id = config["child id"].asInt(); // child will be bind later

    for (Json::ArrayIndex i = 0; i < config["transforms"].size(); i++)
    {
        if (config["transforms"][i]["type"].asString().compare("srt") == 0)
        {
            this->transforms.push_back(ConfigHelpers::loadSRTTransform(config["transforms"][i]));
        }
        else
        {
            throw std::runtime_error("invalid transform type: " +
                                     config["transforms"][i]["type"].asString());
        }
    }
}

void SRTAnimation::setupMotionTransformBlock(void **dst, uint &num_keys, uint &size)
{
    // calc size
    const uint datasize = sizeof(OptixSRTData);
    size = sizeof(OptixSRTMotionTransform) + max(0, num_keys - 2) * datasize;
    // init transform block
    *dst = malloc(size);
    memset(*dst, 0, size);
    // set values
    OptixSRTMotionTransform *motion_transform = reinterpret_cast<OptixSRTMotionTransform *>(*dst);
    motion_transform->child = this->child->handle;
    motion_transform->motionOptions.numKeys = max(2, num_keys);
    motion_transform->motionOptions.flags = this->flags;
    motion_transform->motionOptions.timeBegin = this->tmin;
    motion_transform->motionOptions.timeEnd = this->tmax;
    // copy data
    OptixSRTData *p_data = motion_transform->srtData;
    uint offset = 0;
    for (auto &t : this->transforms)
    {
        p_data[offset] = t->getSRTData();
        offset += 1;
    }
    if (num_keys == 1) // fill 2nd key
    {
        p_data[1] = p_data[0];
    }
}

void MatrixAnimation::postBuild(Scene &sc)
{
    this->linkChild(sc.materials, sc.mediums, sc.meshes, sc.lights, sc.instances);
}

void SRTAnimation::postBuild(Scene &sc)
{
    this->linkChild(sc.materials, sc.mediums, sc.meshes, sc.lights, sc.instances);
}